#!/usr/bin/env python3
"""
管理Milvus集合工具 - 支持列出、删除和清空集合

用法:
  python data_manager.py list               - 列出所有Milvus集合
  python data_manager.py delete <collection> - 删除指定的Milvus集合
  python data_manager.py delete-all         - 删除所有Milvus集合
  python data_manager.py truncate <collection> - 清空指定集合的内容，但保留集合结构
  python data_manager.py truncate-all       - 清空所有集合的内容，但保留集合结构
"""

import sys
import logging
import os
import shutil
from pymilvus import connections, utility, Collection

# 导入统一的日志配置模块
from log_config import setup_logging
from services.config_service import ConfigService
from category import CategoryManager

# 配置日志
setup_logging()
logger = logging.getLogger(__name__)

class MilvusCollectionManager:
    def __init__(self, host='localhost', port='19530'):
        """初始化Milvus连接"""
        self.connection_args = {
            "host": host,
            "port": port
        }
        self._connect()
    
    def _connect(self):
        """连接到Milvus服务"""
        try:
            connections.connect(
                alias="default",
                host=self.connection_args["host"],
                port=self.connection_args["port"]
            )
            logger.info(f"成功连接到Milvus服务: {self.connection_args}")
        except Exception as e:
            logger.error(f"连接Milvus服务失败: {str(e)}")
            raise
    
    def list_collections(self):
        """列出所有Milvus集合"""
        try:
            collections = utility.list_collections()
            logger.info(f"Milvus中的集合数量: {len(collections)}")
            return collections
        except Exception as e:
            logger.error(f"列出Milvus集合失败: {str(e)}")
            return []
    
    def delete_collection(self, collection_name):
        """删除指定的Milvus集合，同时删除对应的category和saved_images子目录"""
        try:
            if not utility.has_collection(collection_name):
                logger.warning(f"集合 '{collection_name}' 不存在，无需删除")
                logger.info(f"开始检查 '{collection_name}' 的 category 和 image 情况 ...")
            
            # 不能删除default集合，除非用户明确要求
            if collection_name == 'default' and '--force' not in sys.argv:
                logger.warning(f"默认集合 '{collection_name}' 受到保护，使用 '--force' 参数来强制删除")
                return False
            
            # 首先删除Milvus集合
            utility.drop_collection(collection_name)
            logger.info(f"成功删除集合: '{collection_name}'")
            
            # 初始化配置服务和分类管理器
            config_service = ConfigService.get_instance()
            
            try:
                # 1. 删除config.ini中对应的category
                category_config = config_service.get_category_config()
                category_manager = CategoryManager(
                    category_dir=category_config['category_dir'],
                    category_file=category_config['category_file']
                )
                
                # 检查并删除对应的分类
                if category_manager.is_category_exist(collection_name):
                    category_manager.delete_category(collection_name)
                    logger.info(f"成功删除分类: '{collection_name}'")
                else:
                    logger.info(f"分类 '{collection_name}' 不存在，无需删除")
                
                # 2. 删除saved_images_dir下对应的子目录
                saved_images_config = config_service.get_saved_images_config()
                saved_images_dir = saved_images_config['saved_images_dir']
                collection_image_dir = os.path.join(saved_images_dir, collection_name)
                
                if os.path.exists(collection_image_dir):
                    shutil.rmtree(collection_image_dir)
                    logger.info(f"成功删除图片目录: '{collection_image_dir}'")
                else:
                    logger.info(f"图片目录 '{collection_image_dir}' 不存在，无需删除")
                
            except Exception as e:
                # 记录非致命错误，但不影响主要功能
                logger.error(f"清理关联资源时出错: {str(e)}")
                
            return True
        except Exception as e:
            logger.error(f"删除集合 '{collection_name}' 失败: {str(e)}")
            return False
    
    def delete_all_collections(self):
        """删除所有Milvus集合"""
        collections = self.list_collections()
        success_count = 0
        
        for collection_name in collections:
            if self.delete_collection(collection_name):
                success_count += 1
        
        logger.info(f"删除操作完成: 成功删除 {success_count}/{len(collections)} 个集合")
        return success_count
        
    def truncate_collection(self, collection_name):
        """清空指定集合的内容，但保留集合结构"""
        try:
            if not utility.has_collection(collection_name):
                logger.warning(f"集合 '{collection_name}' 不存在")
                return False
            
            # 加载集合并清空内容
            collection = Collection(collection_name)
            collection.load()
            collection.truncate()
            collection.release()
            
            logger.info(f"成功清空集合 '{collection_name}' 的内容")
            return True
        except Exception as e:
            logger.error(f"清空集合 '{collection_name}' 内容失败: {str(e)}")
            return False
        
    def truncate_all_collections(self):
        """清空所有集合的内容，但保留集合结构"""
        collections = self.list_collections()
        success_count = 0
        
        for collection_name in collections:
            if self.truncate_collection(collection_name):
                success_count += 1
        
        logger.info(f"清空操作完成: 成功清空 {success_count}/{len(collections)} 个集合的内容")
        return success_count
    
    def close(self):
        """关闭Milvus连接"""
        try:
            connections.disconnect(alias="default")
            logger.info("已关闭Milvus连接")
        except Exception as e:
            logger.error(f"关闭Milvus连接失败: {str(e)}")

if __name__ == "__main__":
    try:
        # 创建Milvus集合管理器实例
        manager = MilvusCollectionManager()
        
        # 解析命令行参数
        if len(sys.argv) > 1:
            if sys.argv[1] == "list":
                # 列出所有集合
                collections = manager.list_collections()
                print("\nMilvus中的集合:")
                for i, coll in enumerate(collections, 1):
                    print(f"{i}. {coll}")
            elif sys.argv[1] == "delete" and len(sys.argv) > 2:
                # 删除指定集合
                collection_name = sys.argv[2]
                manager.delete_collection(collection_name)
            elif sys.argv[1] == "delete-all":
                # 删除所有集合
                confirm = input("警告: 这将删除所有Milvus集合。是否继续? (y/N): ")
                if confirm.lower() == 'y':
                    manager.delete_all_collections()
                else:
                    print("操作已取消")
            elif sys.argv[1] == "truncate" and len(sys.argv) > 2:
                # 清空指定集合的内容
                collection_name = sys.argv[2]
                manager.truncate_collection(collection_name)
            elif sys.argv[1] == "truncate-all":
                # 清空所有集合的内容
                confirm = input("警告: 这将清空所有Milvus集合的内容。是否继续? (y/N): ")
                if confirm.lower() == 'y':
                    manager.truncate_all_collections()
                else:
                    print("操作已取消")
            else:
                print("用法:")
                print("  python data_manager.py list               - 列出所有Milvus集合")
                print("  python data_manager.py delete <collection> - 删除指定的Milvus集合")
                print("  python data_manager.py delete-all         - 删除所有Milvus集合")
                print("  python data_manager.py truncate <collection> - 清空指定集合的内容，但保留集合结构")
                print("  python data_manager.py truncate-all       - 清空所有集合的内容，但保留集合结构")
                print("\n注意: 要删除default集合，请在命令后添加 '--force' 参数")
        else:
            # 默认列出所有集合
            collections = manager.list_collections()
            print("\nMilvus中的集合:")
            for i, coll in enumerate(collections, 1):
                print(f"{i}. {coll}")
            print("\n使用 'python milvus_manager.py -h' 查看更多命令选项")
        
        # 关闭连接
        manager.close()
    except Exception as e:
        logger.error(f"操作失败: {str(e)}")
        sys.exit(1)